"""
WebDriver管理器
负责浏览器驱动的初始化和配置
"""
import os
import time
from selenium import webdriver
from selenium.webdriver.chrome.service import Service as ChromeService
from selenium.webdriver.firefox.service import Service as FirefoxService
from selenium.webdriver.edge.service import Service as EdgeService
from selenium.webdriver.chrome.options import Options as ChromeOptions
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.webdriver.edge.options import Options as EdgeOptions
from webdriver_manager.chrome import ChromeDriverManager
from webdriver_manager.firefox import GeckoDriverManager
from webdriver_manager.microsoft import EdgeChromiumDriverManager
from config.test_config import TestConfig
from utils.logger import Logger


class WebDriverManager:
    """WebDriver管理器类"""
    
    def __init__(self):
        self.logger = Logger().get_logger()
        self.driver = None
    
    def get_driver(self, browser=None):
        """
        获取WebDriver实例
        
        Args:
            browser (str): 浏览器类型 (chrome, firefox, edge)
            
        Returns:
            WebDriver: 浏览器驱动实例
        """
        if browser is None:
            browser = TestConfig.BROWSER
        
        try:
            if browser.lower() == 'chrome':
                self.driver = self._create_chrome_driver()
            elif browser.lower() == 'firefox':
                self.driver = self._create_firefox_driver()
            elif browser.lower() == 'edge':
                self.driver = self._create_edge_driver()
            else:
                raise ValueError(f"不支持的浏览器类型: {browser}")
            
            # 设置窗口大小
            self.driver.set_window_size(
                TestConfig.WINDOW_WIDTH, 
                TestConfig.WINDOW_HEIGHT
            )
            
            # 设置超时时间
            self.driver.implicitly_wait(TestConfig.IMPLICIT_WAIT)
            self.driver.set_page_load_timeout(TestConfig.PAGE_LOAD_TIMEOUT)
            
            self.logger.info(f"成功创建 {browser} 浏览器驱动")
            return self.driver
            
        except Exception as e:
            self.logger.error(f"创建浏览器驱动失败: {str(e)}")
            raise
    
    def _create_chrome_driver(self):
        """创建Chrome浏览器驱动"""
        options = ChromeOptions()
        
        # 添加浏览器选项
        for arg in TestConfig.get_browser_options().get('args', []):
            options.add_argument(arg)
        
        # 设置下载路径
        prefs = {
            "download.default_directory": os.path.join(os.getcwd(), "downloads"),
            "download.prompt_for_download": False,
            "download.directory_upgrade": True,
            "safebrowsing.enabled": True
        }
        options.add_experimental_option("prefs", prefs)
        
        # 创建服务
        service = ChromeService(ChromeDriverManager().install())
        
        return webdriver.Chrome(service=service, options=options)
    
    def _create_firefox_driver(self):
        """创建Firefox浏览器驱动"""
        options = FirefoxOptions()
        
        # 添加浏览器选项
        for arg in TestConfig.get_browser_options().get('args', []):
            options.add_argument(arg)
        
        # 设置下载路径
        options.set_preference("browser.download.folderList", 2)
        options.set_preference("browser.download.manager.showWhenStarting", False)
        options.set_preference("browser.download.dir", os.path.join(os.getcwd(), "downloads"))
        options.set_preference("browser.helperApps.neverAsk.saveToDisk", 
                              "application/pdf,application/x-pdf")
        
        # 创建服务
        service = FirefoxService(GeckoDriverManager().install())
        
        return webdriver.Firefox(service=service, options=options)
    
    def _create_edge_driver(self):
        """创建Edge浏览器驱动"""
        options = EdgeOptions()
        
        # 添加浏览器选项
        for arg in TestConfig.get_browser_options().get('args', []):
            options.add_argument(arg)
        
        # 创建服务
        service = EdgeService(EdgeChromiumDriverManager().install())
        
        return webdriver.Edge(service=service, options=options)
    
    def quit_driver(self):
        """关闭浏览器驱动"""
        if self.driver:
            try:
                self.driver.quit()
                self.logger.info("浏览器驱动已关闭")
            except Exception as e:
                self.logger.error(f"关闭浏览器驱动失败: {str(e)}")
            finally:
                self.driver = None
    
    def take_screenshot(self, filename=None):
        """
        截图
        
        Args:
            filename (str): 文件名，如果不指定则使用时间戳
            
        Returns:
            str: 截图文件路径
        """
        if not self.driver:
            self.logger.error("浏览器驱动未初始化")
            return None
        
        try:
            # 确保截图目录存在
            os.makedirs(TestConfig.SCREENSHOT_DIR, exist_ok=True)
            
            if filename is None:
                timestamp = time.strftime("%Y%m%d_%H%M%S")
                filename = f"screenshot_{timestamp}.png"
            
            filepath = os.path.join(TestConfig.SCREENSHOT_DIR, filename)
            self.driver.save_screenshot(filepath)
            
            self.logger.info(f"截图已保存: {filepath}")
            return filepath
            
        except Exception as e:
            self.logger.error(f"截图失败: {str(e)}")
            return None
    
    def get_page_source(self):
        """获取页面源码"""
        if not self.driver:
            self.logger.error("浏览器驱动未初始化")
            return None
        
        try:
            return self.driver.page_source
        except Exception as e:
            self.logger.error(f"获取页面源码失败: {str(e)}")
            return None
    
    def get_current_url(self):
        """获取当前页面URL"""
        if not self.driver:
            self.logger.error("浏览器驱动未初始化")
            return None
        
        try:
            return self.driver.current_url
        except Exception as e:
            self.logger.error(f"获取当前URL失败: {str(e)}")
            return None
    
    def get_title(self):
        """获取页面标题"""
        if not self.driver:
            self.logger.error("浏览器驱动未初始化")
            return None
        
        try:
            return self.driver.title
        except Exception as e:
            self.logger.error(f"获取页面标题失败: {str(e)}")
            return None 